"""
nnunet/preprocessing/preprocessing.py:
    - preprocess_test_case
        - crop_from_list_of_files(preprocessing/cropping.py)
            - crop
                - crop_to_nonzero
                    - create_nonzero_mask
                        - from scipy.ndimage import binary_fill_holes
                            def create_nonzero_mask(data):
                                from scipy.ndimage import binary_fill_holes
                                assert len(data.shape) == 4 or len(data.shape) == 3, "data must have shape (C, X, Y, Z) or shape (C, X, Y)"
                                nonzero_mask = np.zeros(data.shape[1:], dtype=bool)
                                for c in range(data.shape[0]):
                                    this_mask = data[c] != 0
                                    nonzero_mask = nonzero_mask | this_mask
                                nonzero_mask = binary_fill_holes(nonzero_mask)
                                return nonzero_mask
                        - 此处的功能大致为寻找到图像有内容的区域(像素非黑色)，裁切出来并记录位置参数
        - resample_and_normalize
            - resample_patient
                - resample_data_or_seg，from skimage.transform import resize
                    - if np.any(shape != new_shape)
                    - new_shape = np.round(((np.array(original_spacing) / np.array(target_spacing)).astype(float) * shape)).astype(int)
                    - original_spacing_transposed = np.array(properties["original_spacing"])[self.transpose_forward]
                    - properties["original_spacing"] = np.array(data_itk[0].GetSpacing())[[2, 1, 0]]
                    - 此处大意应是统一3D输入时，xyz大小不统一的问题，target_spacing应是预处理步骤中计算了整体数据得到
                    - 2D可以忽略这个步骤
            -  for c in range(len(data))
                - data[c][mask] = (data[c][mask] - data[c][mask].mean()) / (data[c][mask].std() + 1e-8)
                - mask = seg[-1] >= 0
                - 数据归一化操作如上，seg来自crop, 如果全数据集大小一致，输入是也一致可以忽略
"""
import numpy as np
from scipy.ndimage import binary_fill_holes


def create_nonzero_mask(data):
    assert len(data.shape) == 4 or len(data.shape) == 3, "data must have shape (C, X, Y, Z) or shape (C, X, Y)"
    nonzero_mask = np.zeros(data.shape[1:], dtype=bool)
    for c in range(data.shape[0]):
        this_mask = data[c] != 0
        nonzero_mask = nonzero_mask | this_mask
    nonzero_mask = binary_fill_holes(nonzero_mask)
    return nonzero_mask


def get_bbox_from_mask(mask, outside_value=0):
    mask_voxel_coords = np.where(mask != outside_value)
    minzidx = int(np.min(mask_voxel_coords[0]))
    maxzidx = int(np.max(mask_voxel_coords[0])) + 1
    minxidx = int(np.min(mask_voxel_coords[1]))
    maxxidx = int(np.max(mask_voxel_coords[1])) + 1
    minyidx = int(np.min(mask_voxel_coords[2]))
    maxyidx = int(np.max(mask_voxel_coords[2])) + 1
    return [[minzidx, maxzidx], [minxidx, maxxidx], [minyidx, maxyidx]]


def crop_to_bbox(image, bbox):
    assert len(image.shape) == 3, "only supports 3d images"
    resizer = (slice(bbox[0][0], bbox[0][1]), slice(bbox[1][0], bbox[1][1]), slice(bbox[2][0], bbox[2][1]))
    return image[resizer]


def crop_to_nonzero(data, seg=None, nonzero_label=-1):
    """

    :param data:
    :param seg:
    :param nonzero_label: this will be written into the segmentation map
    :return:
    """
    nonzero_mask = create_nonzero_mask(data)
    bbox = get_bbox_from_mask(nonzero_mask, 0)

    cropped_data = []
    for c in range(data.shape[0]):
        cropped = crop_to_bbox(data[c], bbox)
        cropped_data.append(cropped[None])
    data = np.vstack(cropped_data)

    if seg is not None:
        cropped_seg = []
        for c in range(seg.shape[0]):
            cropped = crop_to_bbox(seg[c], bbox)
            cropped_seg.append(cropped[None])
        seg = np.vstack(cropped_seg)

    nonzero_mask = crop_to_bbox(nonzero_mask, bbox)[None]
    if seg is not None:
        seg[(seg == 0) & (nonzero_mask == 0)] = nonzero_label
    else:
        nonzero_mask = nonzero_mask.astype(int)
        nonzero_mask[nonzero_mask == 0] = nonzero_label
        nonzero_mask[nonzero_mask > 0] = 0
        seg = nonzero_mask
    return data, seg, bbox


def data_normalize(data, seg):
    for c in range(len(data)):
        mask = np.ones(seg.shape[1:], dtype=bool)
        data[c][mask] = (data[c][mask] - data[c][mask].mean()) / (data[c][mask].std() + 1e-8)
        data[c][mask == 0] = 0
    return data


if __name__ == '__main__':
    # binary_fill_holes，此函数需要替换为更加精简的表达，否则难以完成cpp重写
    import cv2

    mat = cv2.imread("../resource/raw_resize/2.png")
    mat = mat.transpose((2, 0, 1))
    mat = np.expand_dims(mat, 1)
    mat = mat.astype(np.float32)
    print("Mat Shape", mat.shape)

    data, seg, _ = crop_to_nonzero(mat)
    out_data = data_normalize(data, seg)
    print(out_data.shape)
